package com.hqq.search.union_find;

/**
 * WeightQuickUnion
 * 优化的QuickUnion
 * 目的:
 * 优化树的高度 防止出现单支树影响查找性能
 * 时间复杂度:
 * LogN
 * 实现方法:
 * 使用数组保存节点的子节点个数 合并节点时进行判断 将子节点小的树插入到子节点大的树下
 * Created by heqianqian on 2017/8/3.
 */
public class WeightQuickUnion implements UnionApi{

    private int[] idArray;

    /**
     * 每个节点的子节点个数
     * 对树进行优化
     */
    private int[] nodeSize;

    /**
     * 节点个数
     */
    private int count;

    public WeightQuickUnion(int count) {
        this.count = count;
        this.idArray = new int[count];
        this.nodeSize = new int[count];
        for (int i = 0; i < count; i++) {
            idArray[i] = i;
            nodeSize[i] = 1;//初始为1
        }
    }

    /**
     * 获得P节点的根节点
     */
    public int root(int p) {
        while (p != idArray[p]) {
            p = idArray[p];
        }
        return p;
    }

    @Override
    public void union(int p, int q) {
        int pRootIndex = root(p);
        int qRootIndex = root(q);
        if (qRootIndex == pRootIndex) {
            return;
        }
        //将节点个数小的树连接到节点个数大的树上 防止树的高度过高 搜索效率低
        if (nodeSize[pRootIndex] <= nodeSize[qRootIndex]) {
            //把p连接到q上
            idArray[pRootIndex] = qRootIndex;
            nodeSize[qRootIndex] += nodeSize[pRootIndex];
        } else {
            idArray[qRootIndex] = pRootIndex;
            nodeSize[pRootIndex] += nodeSize[qRootIndex];
        }
        count--;
    }

    @Override
    public boolean isConnected(int p, int q) {
        return root(p) == root(q);
    }

    @Override
    public int find(int p) {
        return idArray[p];
    }

    @Override
    public int count() {
        return this.count;
    }

    @Override
    public int[] getArray() {
        return idArray;
    }

    public int[] getNodeSize() {
        return nodeSize;
    }
}
